import asyncio
from src.module.Model import Model
from pydantic import BaseModel, Field
from langchain_core.tools import tool, StructuredTool, ToolException
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_community.tools.tavily_search import TavilySearchResults

"""
    :descript: 普通的tools，一定要写tools的描述, 可以是个异步tools也可以是个普通tools
"""
@tool
async def ordinary(a: int, b: int) -> int:
    """ 将这两个数相加 """
    return a + b

class CalculatorInput(BaseModel):
    a: int = Field(description='第一个参数')
    b: int = Field(description='第二个参数')

class SeniorTools(Model):
    @staticmethod
    def add_tools(a: int, b: int) -> list:
        """ 将这两个数据添加到新数组中 """
        raise ToolException(123)

    @staticmethod
    async def del_tools(a: int, b: int) -> int:
        """ 返回第一个数减去第二个数的值 """
        return a - b

    """
        :descript 执行 工具内部的自定义 tool
    """
    async def exec(self):
        tools_func = StructuredTool.from_function(
            func=self.add_tools,    ## 指定同步运行的tools函数
            coroutine=self.del_tools, ## 指定运行的异步tools函数
            name='struc_tools', ## 指定 structured tool 的名称
            description='自定义tools', ## 指定 structured tool 的描述
            args_schema=CalculatorInput, ## 自定 tools 的参数值, 同时也说明了，一个 structured tool 使用的 tool 函数只能是固定的参数
            return_direct=False, ## 仅对代理相关。当为True时，在调⽤给定⼯具后，代理将停⽌并将结果直接返回给⽤户。
            handle_tool_error='全局错误拦截' ## 工具全局错误拦截, 只能拦截：ToolException 的错误类型，不能拦截其它的错误类型
        )

        print(tools_func.invoke({ 'a': 1, 'b': 2 }))
        print(await tools_func.ainvoke({'a': 6, 'b': 0}))

    """
        :description 使用维基百科的 tool
    """
    @staticmethod
    def wikipedia():
        api_wrapper = WikipediaAPIWrapper(top_k_result=1)
        wikipedia_tool = WikipediaQueryRun(api_wrapper = api_wrapper)
        print(wikipedia_tool.invoke({ 'query': '猫' }))

    """
        :description tools 不单单可以作为工具，同时也能作为一个智能体使用
    """
    @staticmethod
    def tvly_search():
        search = TavilySearchResults(max_results=1)
        print(search.invoke('深圳今天的天气怎么样'))

    def start(self):
        # asyncio.run(self.exec())
        # self.wikipedia()
        self.tvly_search()
